import time
import threading
from flask import Flask, Response
from PIL import Image
import board, digitalio
from adafruit_rgb_display import st7735
from picamera2 import Picamera2
import cv2
import numpy as np

# ========= 1) 初始化 TFT 屏幕 =========
spi = board.SPI()
cs = digitalio.DigitalInOut(board.D23)    # 接线：CS -> GPIO23
dc = digitalio.DigitalInOut(board.D24)    # DC -> GPIO24
rst = digitalio.DigitalInOut(board.D25)   # RST -> GPIO25

# 初始化 ST7735R 屏幕（128x160 竖屏）
display = st7735.ST7735R(
    spi,
    cs=cs,
    dc=dc,
    rst=rst,
    baudrate=24000000,
    width=128,
    height=160,
    rotation=0,
    bgr=True  # 屏幕需要 BGR 格式
)
PANEL_W, PANEL_H = display.width, display.height
print(f"✅ TFT 屏幕初始化完成：{PANEL_W}x{PANEL_H}")

# 屏幕自检（红绿蓝）
for color in [(255,0,0), (0,255,0), (0,0,255)]:
    img = Image.new("RGB", (PANEL_W, PANEL_H), color)
    display.image(img)
    time.sleep(0.3)

# ========= 2) 初始化摄像头 =========
picam2 = Picamera2()
# 配置为 RGB 格式（同时满足网页和屏幕需求）
config = picam2.create_video_configuration(
    main={"size": (640, 480), "format": "RGB888"}  # 640x480 兼顾流畅度和性能
)
picam2.configure(config)
picam2.start()
print("✅ 摄像头启动成功")

# 全局变量：存储最新一帧图像（供多线程共享）
latest_frame = None
frame_lock = threading.Lock()  # 线程锁，避免资源竞争

# ========= 3) TFT 屏幕显示线程 =========
def tft_display_loop():
    """单独线程：持续刷新TFT屏幕（修复颜色颠倒）"""
    global latest_frame
    while True:
        with frame_lock:
            # 复制最新一帧（避免被主线程修改）
            frame = latest_frame.copy() if latest_frame is not None else None
        
        if frame is not None:
            # 处理图像适配TFT屏幕
            # 1. 旋转90度（将横屏640x480转为竖屏480x640，适配TFT的竖屏）
            frame_rotated = cv2.rotate(frame, cv2.ROTATE_90_CLOCKWISE)
            # 2. 缩放到TFT尺寸（128x160）
            frame_scaled = cv2.resize(frame_rotated, (PANEL_W, PANEL_H), interpolation=cv2.INTER_LINEAR)
            # 3. 关键修复：将RGB转为BGR（适配屏幕的bgr=True设置）
            frame_bgr = cv2.cvtColor(frame_scaled, cv2.COLOR_RGB2BGR)
            # 4. 转为PIL图像并显示（此时传入的是BGR格式，与屏幕要求一致）
            img = Image.fromarray(frame_bgr)
            display.image(img)
        
        # 控制屏幕刷新频率（约30fps）
        time.sleep(0.03)

# ========= 4) Flask 网页流媒体服务 =========
app = Flask(__name__)

def generate_mjpeg():
    """生成MJPEG流，供网页播放"""
    global latest_frame
    while True:
        # 读取摄像头帧（RGB格式）
        frame = picam2.capture_array()
        # 用线程锁保护全局变量
        with frame_lock:
            latest_frame = frame.copy()
        
        # 直接用RGB格式编码为JPEG（网页需要RGB）
        ret, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 80])
        if not ret:
            continue
        
        # 按MJPEG协议输出
        yield (b'--frame\r\n'
               b'Content-Type: image/jpeg\r\n\r\n' + buffer.tobytes() + b'\r\n')

@app.route('/video_feed')
def video_feed():
    return Response(
        generate_mjpeg(),
        mimetype='multipart/x-mixed-replace; boundary=frame'
    )

@app.route('/')
def index():
    return '''
    <html>
      <head><title>双屏显示 - 网页+TFT</title></head>
      <body>
        <h2>网页实时画面</h2>
        <img src="/video_feed" width="640" height="480">
      </body>
    </html>
    '''

# ========= 5) 启动程序 =========
if __name__ == '__main__':
    # 启动TFT显示线程
    tft_thread = threading.Thread(target=tft_display_loop, daemon=True)
    tft_thread.start()
    print("✅ TFT显示线程启动")
    
    # 启动Flask服务（支持多线程访问）
    app.run(host='0.0.0.0', port=5000, threaded=True)